package par

import (
	"fmt"
	"strconv"

	"gitcode.com/deyiyangyang/bampoo/ast"
	"gitcode.com/deyiyangyang/bampoo/lex"

	"gitcode.com/deyiyangyang/bampoo/tok"
)

// Parser represents a Scheme parser
type Parser struct {
	lexer     *lex.Lexer
	curToken  tok.Token
	peekToken tok.Token
	errors    []string
}

// New creates a new Parser
func New(lexer *lex.Lexer) *Parser {
	p := &Parser{
		lexer:  lexer,
		errors: []string{},
	}

	// Read two tokens, so curToken and peekToken are both set
	p.nextToken()
	p.nextToken()

	return p
}

// Errors returns parsing errors
func (p *Parser) Errors() []string {
	return p.errors
}

// CurToken returns current token (for debugging)
func (p *Parser) CurToken() tok.Token {
	return p.curToken
}

// PeekToken returns peek token (for debugging)
func (p *Parser) PeekToken() tok.Token {
	return p.peekToken
}

func (p *Parser) nextToken() {
	p.curToken = p.peekToken
	p.peekToken = p.lexer.NextToken()
}

// ParseProgram parses the entire program
func (p *Parser) ParseProgram() *ast.Program {
	program := &ast.Program{}
	program.Expressions = []ast.Expression{}

	for p.curToken.Type != tok.EOF {
		expr := p.parseExpression()
		if expr != nil {
			program.Expressions = append(program.Expressions, expr)
		}
		// parseExpression already advanced curToken to the next token after the expression
		// so we don't need to advance it here, just continue the loop
	}

	return program
}

func (p *Parser) parseExpression() ast.Expression {
	switch p.curToken.Type {
	case tok.LPAREN:
		return p.parseList()
	case tok.STRING:
		lit := &ast.StringLiteral{Value: p.curToken.Literal}
		p.nextToken() // consume string literal
		return lit
	case tok.INT:
		value, err := strconv.ParseInt(p.curToken.Literal, 10, 64)
		if err != nil {
			p.errors = append(p.errors, fmt.Sprintf("could not parse %q as integer", p.curToken.Literal))
			return nil
		}
		lit := &ast.IntegerLiteral{Value: value}
		p.nextToken() // consume integer literal
		return lit
	case tok.REAL:
		value, err := strconv.ParseFloat(p.curToken.Literal, 64)
		if err != nil {
			p.errors = append(p.errors, fmt.Sprintf("could not parse %q as real number", p.curToken.Literal))
			return nil
		}
		lit := &ast.RealLiteral{Value: value}
		p.nextToken() // consume real literal
		return lit
	case tok.BOOLEAN:
		var value bool
		if p.curToken.Literal == "#t" {
			value = true
		} else {
			value = false
		}
		lit := &ast.BooleanLiteral{Value: value}
		p.nextToken() // consume boolean literal
		return lit
	case tok.SYMBOL:
		sym := &ast.Symbol{Value: p.curToken.Literal}
		p.nextToken() // consume symbol
		return sym
	case tok.QUOTE:
		p.nextToken() // consume quote token
		expr := p.parseExpression()
		if expr == nil {
			p.errors = append(p.errors, "expected expression after quote")
			return nil
		}
		return &ast.QuoteExpr{Expression: expr}
	case tok.DEFINE, tok.DISPLAY:
		// These should only appear inside lists, not at top level
		p.errors = append(p.errors, fmt.Sprintf("unexpected token at top level: %s", p.curToken.Type))
		p.nextToken() // consume the token
		return nil
	default:
		p.errors = append(p.errors, fmt.Sprintf("unexpected token: %s", p.curToken.Type))
		p.nextToken() // consume the unexpected token to avoid infinite loop
		return nil
	}
}

func (p *Parser) parseList() ast.Expression {
	if p.curToken.Type != tok.LPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected '(', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume '('

	if p.curToken.Type == tok.RPAREN {
		p.nextToken() // consume ')'
		return &ast.List{Elements: []ast.Expression{}}
	}

	// Check if this is a special form
	switch p.curToken.Type {
	case tok.DEFINE:
		return p.parseDefine()
	case tok.SET:
		return p.parseSet()
	case tok.DISPLAY:
		return p.parseDisplay()
	case tok.BEGIN:
		return p.parseBegin()
	case tok.LAMBDA:
		return p.parseLambda()
	case tok.LET:
		return p.parseLet()
	case tok.IF:
		return p.parseIf()
	case tok.COND:
		return p.parseCond()
	case tok.AND, tok.OR, tok.NOT:
		return p.parseLogical()
	case tok.PLUS, tok.MINUS, tok.MULTIPLY, tok.DIVIDE:
		return p.parseArithmetic()
	case tok.GREATER, tok.LESS, tok.EQUAL, tok.GREATER_OR_EQUAL, tok.LESS_OR_EQUAL:
		return p.parseComparison()
	case tok.SYMBOL:
		// Check for special form with symbol name
		if p.curToken.Literal == "quote" {
			return p.parseQuote()
		}
		return p.parseFunctionCall()
	default:
		return p.parseFunctionCall()
	}
}

func (p *Parser) parseDefine() ast.Expression {
	if p.curToken.Type != tok.DEFINE {
		p.errors = append(p.errors, fmt.Sprintf("expected 'define', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume 'define'

	// Check if this is a function definition: (define (name params...) body...)
	if p.curToken.Type == tok.LPAREN {
		return p.parseFunctionDefine()
	}

	// Regular variable definition
	if p.curToken.Type != tok.SYMBOL {
		p.errors = append(p.errors, fmt.Sprintf("expected symbol after 'define', got %s", p.curToken.Type))
		return nil
	}

	name := ast.Symbol{Value: p.curToken.Literal}
	p.nextToken() // consume symbol name

	value := p.parseExpression()
	if value == nil {
		p.errors = append(p.errors, "expected value after define name")
		return nil
	}

	// After parsing the value expression, curToken should be the token after the value
	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.DefineExpr{
		Name:  name,
		Value: value,
	}
}

func (p *Parser) parseQuote() ast.Expression {
	if p.curToken.Type != tok.SYMBOL || p.curToken.Literal != "quote" {
		p.errors = append(p.errors, fmt.Sprintf("expected 'quote', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume 'quote'

	expr := p.parseExpression()
	if expr == nil {
		p.errors = append(p.errors, "expected expression after 'quote'")
		return nil
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.QuoteExpr{Expression: expr}

}

func (p *Parser) parseFunctionDefine() ast.Expression {
	p.nextToken() // consume '('

	if p.curToken.Type != tok.SYMBOL {
		p.errors = append(p.errors, fmt.Sprintf("expected function name, got %s", p.curToken.Type))
		return nil
	}

	name := ast.Symbol{Value: p.curToken.Literal}
	p.nextToken() // consume function name

	// Parse parameters
	parameters := []ast.Symbol{}
	for p.curToken.Type != tok.RPAREN && p.curToken.Type != tok.EOF {
		if p.curToken.Type != tok.SYMBOL {
			p.errors = append(p.errors, fmt.Sprintf("expected parameter name, got %s", p.curToken.Type))
			return nil
		}
		parameters = append(parameters, ast.Symbol{Value: p.curToken.Literal})
		p.nextToken() // consume parameter name
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	// Parse body
	body := []ast.Expression{}
	for p.curToken.Type != tok.RPAREN && p.curToken.Type != tok.EOF {
		expr := p.parseExpression()
		if expr != nil {
			body = append(body, expr)
		}
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	// Create a lambda expression for the function body
	lambda := &ast.LambdaExpr{
		Parameters: parameters,
		Body:       body,
	}

	return &ast.DefineExpr{
		Name:  name,
		Value: lambda,
	}
}

func (p *Parser) parseDisplay() ast.Expression {
	if p.curToken.Type != tok.DISPLAY {
		p.errors = append(p.errors, fmt.Sprintf("expected 'display', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume 'display'

	value := p.parseExpression()
	if value == nil {
		p.errors = append(p.errors, "expected value after 'display'")
		return nil
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.DisplayExpr{
		Value: value,
	}
}

func (p *Parser) parseSet() ast.Expression {
	if p.curToken.Type != tok.SET {
		p.errors = append(p.errors, fmt.Sprintf("expected 'set!', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume 'set!'

	// Parse variable name
	if p.curToken.Type != tok.SYMBOL {
		p.errors = append(p.errors, fmt.Sprintf("expected symbol after 'set!', got %s", p.curToken.Type))
		return nil
	}

	name := ast.Symbol{Value: p.curToken.Literal}
	p.nextToken() // consume symbol name

	// Parse new value
	value := p.parseExpression()
	if value == nil {
		p.errors = append(p.errors, "expected value after set! variable name")
		return nil
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.SetExpr{
		Name:  name,
		Value: value,
	}
}

// parseIf parses an if special form
// Syntax: (if condition consequence [alternative])
func (p *Parser) parseIf() ast.Expression {
	if p.curToken.Type != tok.IF {
		p.errors = append(p.errors, fmt.Sprintf("expected 'if', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume 'if'

	condition := p.parseExpression()
	consequence := p.parseExpression()
	var alternative ast.Expression

	// Check if there's an alternative clause
	if p.curToken.Type != tok.RPAREN {
		alternative = p.parseExpression()
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.IfExpr{
		Condition:   condition,
		Consequence: consequence,
		Alternative: alternative,
	}
}

// parseCond parses a cond special form
// Syntax: (cond (test1 expr1) (test2 expr2) ... (else exprN))
func (p *Parser) parseCond() ast.Expression {
	if p.curToken.Type != tok.COND {
		p.errors = append(p.errors, fmt.Sprintf("expected 'cond', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume 'cond'

	var clauses []ast.CondClause

	for p.curToken.Type != tok.RPAREN && p.curToken.Type != tok.EOF {
		if p.curToken.Type != tok.LPAREN {
			p.errors = append(p.errors, "Expected '(' at start of cond clause")
			return nil
		}

		p.nextToken() // consume '('

		clause := ast.CondClause{}

		// Check if this is an else clause
		if p.curToken.Type == tok.SYMBOL && p.curToken.Literal == "else" {
			clause.IsElse = true
			p.nextToken() // consume 'else'
			clause.Result = p.parseExpression()
		} else {
			// Regular clause
			clause.Test = p.parseExpression()
			clause.Result = p.parseExpression()
		}

		if p.curToken.Type != tok.RPAREN {
			p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
			return nil
		}

		p.nextToken() // consume ')'
		clauses = append(clauses, clause)
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.CondExpr{Clauses: clauses}
}

// parseComparison parses a comparison expression
// Syntax: (> expr1 expr2), (< expr1 expr2), etc.
func (p *Parser) parseComparison() ast.Expression {
	operator := p.curToken.Literal
	p.nextToken() // consume operator

	var operands []ast.Expression
	for p.curToken.Type != tok.RPAREN && p.curToken.Type != tok.EOF {
		operand := p.parseExpression()
		if operand != nil {
			operands = append(operands, operand)
		}
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.ComparisonExpr{Operator: operator, Operands: operands}
}

// parseLogical parses a logical expression
// Syntax: (and expr1 expr2 ...), (or expr1 expr2 ...), (not expr)
func (p *Parser) parseLogical() ast.Expression {
	operator := p.curToken.Literal
	p.nextToken() // consume operator

	var operands []ast.Expression
	for p.curToken.Type != tok.RPAREN && p.curToken.Type != tok.EOF {
		operand := p.parseExpression()
		if operand != nil {
			operands = append(operands, operand)
		}
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.LogicalExpr{Operator: operator, Operands: operands}
}

func (p *Parser) parseArithmetic() ast.Expression {
	operator := p.curToken.Literal
	p.nextToken() // consume operator

	var operands []ast.Expression
	for p.curToken.Type != tok.RPAREN && p.curToken.Type != tok.EOF {
		operand := p.parseExpression()
		if operand != nil {
			operands = append(operands, operand)
		}
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.ArithmeticExpr{Operator: operator, Operands: operands}
}

func (p *Parser) parseLambda() ast.Expression {
	if p.curToken.Type != tok.LAMBDA {
		p.errors = append(p.errors, fmt.Sprintf("expected 'lambda', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume 'lambda'

	if p.curToken.Type != tok.LPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected '(' after 'lambda', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume '('

	// Parse parameters
	parameters := []ast.Symbol{}
	for p.curToken.Type != tok.RPAREN && p.curToken.Type != tok.EOF {
		if p.curToken.Type != tok.SYMBOL {
			p.errors = append(p.errors, fmt.Sprintf("expected parameter name, got %s", p.curToken.Type))
			return nil
		}
		parameters = append(parameters, ast.Symbol{Value: p.curToken.Literal})
		p.nextToken() // consume parameter name
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	// Parse body
	body := []ast.Expression{}
	for p.curToken.Type != tok.RPAREN && p.curToken.Type != tok.EOF {
		expr := p.parseExpression()
		if expr != nil {
			body = append(body, expr)
		}
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.LambdaExpr{
		Parameters: parameters,
		Body:       body,
	}
}

func (p *Parser) parseBegin() ast.Expression {
	if p.curToken.Type != tok.BEGIN {
		p.errors = append(p.errors, fmt.Sprintf("expected 'begin', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume 'begin'

	// Parse body expressions
	body := []ast.Expression{}
	for p.curToken.Type != tok.RPAREN && p.curToken.Type != tok.EOF {
		expr := p.parseExpression()
		if expr != nil {
			body = append(body, expr)
		}
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.BeginExpr{
		Body: body,
	}
}

func (p *Parser) parseLet() ast.Expression {
	if p.curToken.Type != tok.LET {
		p.errors = append(p.errors, fmt.Sprintf("expected 'let', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume 'let'

	if p.curToken.Type != tok.LPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected '(' after 'let', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume '('

	// Parse bindings
	bindings := []ast.LetBinding{}
	for p.curToken.Type != tok.RPAREN && p.curToken.Type != tok.EOF {
		if p.curToken.Type != tok.LPAREN {
			p.errors = append(p.errors, fmt.Sprintf("expected '(' for binding, got %s", p.curToken.Type))
			return nil
		}

		p.nextToken() // consume '('

		if p.curToken.Type != tok.SYMBOL {
			p.errors = append(p.errors, fmt.Sprintf("expected symbol in binding, got %s", p.curToken.Type))
			return nil
		}

		name := ast.Symbol{Value: p.curToken.Literal}
		p.nextToken() // consume symbol name

		value := p.parseExpression()
		if value == nil {
			p.errors = append(p.errors, "expected value in binding")
			return nil
		}

		if p.curToken.Type != tok.RPAREN {
			p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
			return nil
		}

		p.nextToken() // consume ')'

		bindings = append(bindings, ast.LetBinding{
			Name:  name,
			Value: value,
		})
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	// Parse body expressions
	body := []ast.Expression{}
	for p.curToken.Type != tok.RPAREN && p.curToken.Type != tok.EOF {
		expr := p.parseExpression()
		if expr != nil {
			body = append(body, expr)
		}
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.LetExpr{
		Bindings: bindings,
		Body:     body,
	}
}

func (p *Parser) parseFunctionCall() ast.Expression {
	// Parse function
	function := p.parseExpression()
	if function == nil {
		return nil
	}

	// Parse arguments
	arguments := []ast.Expression{}
	for p.curToken.Type != tok.RPAREN && p.curToken.Type != tok.EOF {
		arg := p.parseExpression()
		if arg != nil {
			arguments = append(arguments, arg)
		}
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.FunctionCall{
		Function:  function,
		Arguments: arguments,
	}
}

func (p *Parser) parseRegularList() ast.Expression {
	elements := []ast.Expression{}

	// Parse first element
	elem := p.parseExpression()
	if elem != nil {
		elements = append(elements, elem)
	}

	// Parse remaining elements
	for p.curToken.Type != tok.RPAREN && p.curToken.Type != tok.EOF {
		elem := p.parseExpression()
		if elem != nil {
			elements = append(elements, elem)
		}
	}

	if p.curToken.Type != tok.RPAREN {
		p.errors = append(p.errors, fmt.Sprintf("expected ')', got %s", p.curToken.Type))
		return nil
	}

	p.nextToken() // consume ')'

	return &ast.List{Elements: elements}
}

func (p *Parser) isArithmeticOperator(tokenType tok.TokenType) bool {
	return tokenType == tok.PLUS || tokenType == tok.MINUS || tokenType == tok.MULTIPLY || tokenType == tok.DIVIDE
}

func (p *Parser) isComparisonOperator(tokenType tok.TokenType) bool {
	return tokenType == tok.GREATER || tokenType == tok.LESS || tokenType == tok.EQUAL || 
		tokenType == tok.GREATER_OR_EQUAL || tokenType == tok.LESS_OR_EQUAL
}
